#!/usr/bin/env python3
import argparse
import os
import sys
from time import sleep

# Import P4Runtime lib
sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), '../../utils/'))
import p4runtime_lib.bmv2
import p4runtime_lib.helper
from p4runtime_lib.error_utils import printGrpcError
from p4runtime_lib.switch import ShutdownAllSwitchConnections

def configure_ecn_threshold(p4info_helper, switch, threshold):
    """配置ECN阈值 - 简化版本"""
    try:
        # 使用简单的直接配置，不使用复杂的表匹配
        table_entry = p4info_helper.buildTableEntry(
            table_name="MyIngress.ecn_config",
            match_fields={
                "standard_metadata.ingress_port": 0
            },
            action_name="MyIngress.set_ecn_threshold",
            action_params={
                "threshold": threshold
            })

        switch.WriteTableEntry(table_entry)
        print(f"Successfully set ECN threshold to {threshold} on {switch.name}")
        return True

    except Exception as e:
        print(f"Error setting ECN threshold on {switch.name}: {e}")
        return False

def main(p4info_file_path, bmv2_file_path):
    # 确保日志目录存在
    os.makedirs("logs", exist_ok=True)

    p4info_helper = p4runtime_lib.helper.P4InfoHelper(p4info_file_path)

    try:
        # 创建交换机连接
        s1 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s1', address='127.0.0.1:50051', device_id=0,
            proto_dump_file='logs/s1-p4runtime-requests.txt')

        s2 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s2', address='127.0.0.1:50052', device_id=1,
            proto_dump_file='logs/s2-p4runtime-requests.txt')

        s3 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s3', address='127.0.0.1:50053', device_id=2,
            proto_dump_file='logs/s3-p4runtime-requests.txt')

        switches = [s1, s2, s3]

        # 发送主仲裁更新
        for sw in switches:
            sw.MasterArbitrationUpdate()

        # 安装P4程序
        for sw in switches:
            sw.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                          bmv2_json_file_path=bmv2_file_path)
            print(f"Installed P4 program on {sw.name}")

        # 配置ECN阈值 - 只配置基本的阈值表
        print("\n----- Configuring ECN thresholds -----")
        success_count = 0
        for i, sw in enumerate(switches):
            # 为不同交换机设置不同的阈值
            threshold = [5, 10, 15][i]  # s1:5, s2:10, s3:15
            if configure_ecn_threshold(p4info_helper, sw, threshold):
                success_count += 1

        print(f"\nSuccessfully configured {success_count}/3 switches")

        if success_count == 3:
            print("ECN configuration completed successfully!")
            print("The switches will now use dynamic thresholds for ECN marking.")
            print("Press Ctrl+C to stop the controller...")

            # 保持运行
            while True:
                sleep(10)
        else:
            print("Some switches failed to configure. Check the errors above.")

    except KeyboardInterrupt:
        print("\nShutting down controller.")
    except Exception as e:
        print(f"Error: {e}")
        import traceback
        traceback.print_exc()

    ShutdownAllSwitchConnections()

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='ECN P4Runtime Controller')
    parser.add_argument('--p4info', help='p4info file', default='./build/ecn.p4.p4info.txt')
    parser.add_argument('--bmv2-json', help='BMv2 JSON file', default='./build/ecn.json')
    args = parser.parse_args()

    if not os.path.exists(args.p4info):
        print(f"p4info file not found: {args.p4info}")
        # 尝试其他可能的位置
        args.p4info = './build/ecn.p4info.txt'
        if not os.path.exists(args.p4info):
            print(f"p4info file not found: {args.p4info}")
            exit(1)

    if not os.path.exists(args.bmv2_json):
        print(f"BMv2 JSON file not found: {args.bmv2_json}")
        exit(1)

    main(args.p4info, args.bmv2_json)
